import jax
import jax.numpy as jnp
from jax.scipy.integrate import trapezoid as jsci_trapezoid

__all__ = [
	"calcu_psd_Hc", "calcu_psd_Vns", "calcu_psd_Vns_Hc",
	"calcu_robust_infide_Hc", "calcu_robust_infide_Vns", "calcu_robust_infide_Vns_Hc"
]

def calcu_psd_Hc(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots, ws = [w_1, w_2, ...]

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)
):
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# Hc_ts = jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	Hc_ts = jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + Hc_ts  # shape = (Nt, dim, dim)
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vv = jax.vmap(if_trans, (0, 0), 0)
	if_Hc_ts = if_trans_vv(U_ts, Hc_ts)  # (Nt, dim, dim)
	if_Hc_ws = F_trans(if_Hc_ts)
	F_Hc_ws = jnp.real(jnp.einsum("wde,wde->w", P @ if_Hc_ws, jnp.conj(if_Hc_ws)) / Tr_P)
	return UT, F_Hc_ws

def calcu_robust_infide_Hc(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots, ws = [w_1, w_2, ...]
		# dws,  # diff(ws)

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)
		S_Hc_ws,  # PSD of noises
):
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# Hc_ts = jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	Hc_ts = jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + Hc_ts  # shape = (Nt, dim, dim)
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vv = jax.vmap(if_trans, (0, 0), 0)
	if_Hc_ts = if_trans_vv(U_ts, Hc_ts)  # (Nt, dim, dim)
	if_Hc_ws = F_trans(if_Hc_ts)
	F_Hc_ws = jnp.einsum("wde,wde->w", P @ if_Hc_ws, jnp.conj(if_Hc_ws)) / Tr_P
	robust_infide_Hc = jnp.real(jsci_trapezoid(F_Hc_ws * S_Hc_ws, ws)) / jnp.pi
	return UT, robust_infide_Hc

def calcu_psd_Vns(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)

		Vns,  # noise Hamiltonians, [Vn_1, Vn_2, ..., Vn_m],
):
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# H_tot_ts = H0[None, :] + jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # shape = (Nt, dim, dim)
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vc = jax.vmap(if_trans, (0, None), 0)

	def func_psd_Vn(Vn):
		if_Vn_ts = if_trans_vc(U_ts, Vn)  # (Nt, dim, dim)
		if_Vn_ws = F_trans(if_Vn_ts)
		F_Vn_ws = jnp.real(jnp.einsum("wde,wde->w", P @ if_Vn_ws, jnp.conj(if_Vn_ws)) / Tr_P)
		return F_Vn_ws

	return UT, jax.vmap(func_psd_Vn, 0, 0)(Vns)

def calcu_robust_infide_Vns(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots
		# dws,  # diff(ws)

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)

		Vns,  # noise Hamiltonians, [Vn_1, Vn_2, ..., Vn_m],
		S_Vns_ws,  # PSD of noises
):
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# H_tot_ts = H0[None, :] + jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # avoid 'ComplexWarning' warning
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vc = jax.vmap(if_trans, (0, None), 0)

	def func_robust_infide_Vn(Vn, S_Vn_ws):
		if_Vn_ts = if_trans_vc(U_ts, Vn)  # (Nt, dim, dim)
		if_Vn_ws = F_trans(if_Vn_ts)
		F_Vn_ws = jnp.einsum("wde,wde->w", P @ if_Vn_ws, jnp.conj(if_Vn_ws)) / Tr_P
		# robust_infide_Vn = jnp.real((F_Vn_ws * S_Vn_ws)[:-1] @ dws / jnp.pi)
		# return robust_infide_Vn
		return jnp.real(jsci_trapezoid(F_Vn_ws * S_Vn_ws, ws)) / jnp.pi

	return UT, jax.vmap(func_robust_infide_Vn, (0, 0), 0)(Vns, S_Vns_ws)

def calcu_psd_Vns_Hc(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots, ws = [w_1, w_2, ...]
		# dws,  # diff(ws)

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)

		Vns,  # noise Hamiltonians, [Vn_1, Vn_2, ..., Vn_m],
):
	N_Vns = Vns.shape[0]
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# Hc_ts = jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	Hc_ts = jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + Hc_ts  # shape = (Nt, dim, dim)
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vc = jax.vmap(if_trans, (0, None), 0)
	if_trans_vv = jax.vmap(if_trans, (0, 0), 0)

	def func_psd_from_fw(if_H_ws):
		F_H_ws = jnp.real(jnp.einsum("wde,wde->w", P @ if_H_ws, jnp.conj(if_H_ws)) / Tr_P)
		return F_H_ws

	def func_psd_Vn(Vn):
		if_Vn_ts = if_trans_vc(U_ts, Vn)  # (Nt, dim, dim)
		if_Vn_ws = F_trans(if_Vn_ts)
		return func_psd_from_fw(if_Vn_ws)

	func_psd_Vns = jax.vmap(func_psd_Vn, 0, 0)

	Fs_ws = jnp.zeros((N_Vns + 1, len(ws)))
	Fs_ws = Fs_ws.at[:N_Vns, :].set(func_psd_Vns(Vns))

	if_Hc_ts = if_trans_vv(U_ts, Hc_ts)  # (Nt, dim, dim)
	if_Hc_ws = F_trans(if_Hc_ts)
	Fs_ws = Fs_ws.at[N_Vns, :].set(func_psd_from_fw(if_Hc_ws))
	return UT, Fs_ws

def calcu_robust_infide_Vns_Hc(
		ucs_ts,  # control pulses in time sequences, shape = (n:=len(ucs), N:=len(ts)-1)
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		dts,  # diff(ts)
		ws,  # frequency slots, ws = [w_1, w_2, ...]
		# dws,  # diff(ws)

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)
		S_Hc_ws,  # PSD of noises

		Vns,  # noise Hamiltonians, [Vn_1, Vn_2, ..., Vn_m],
		S_Vns_ws,  # PSD of noises
):
	N_Vns = Vns.shape[0]
	I = jnp.eye(dim, dtype=complex)
	Tr_P = jnp.trace(P)  # trace of the projection matrix

	# 预计算所有时间步的演化矩阵 --------------------------------------------------
	f_ts = 0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :]  # shape = (N_ctrl, Nt)
	# Hc_ts = jnp.einsum("nt,nde->tde", f_ts, Hcs)  # shape = (Nt, dim, dim)
	Hc_ts = jnp.einsum("nt,nde->tde", f_ts.astype(complex), Hcs)  # shape = (Nt, dim, dim)
	H_tot_ts = H0[None, :] + Hc_ts  # shape = (Nt, dim, dim)
	Udt_ts = jax.vmap(jax.scipy.linalg.expm)(-1.j * H_tot_ts)  # shape = (Nt, dim, dim)
	UT, U_ts = jax.lax.scan(lambda U_prev, Udt_t: (Udt_t @ U_prev, U_prev), init=I, xs=Udt_ts)

	# 噪声哈密顿量的傅里叶变换 ------------------------------------------
	phases = jnp.exp(1.j * ws[:, None] * ts[None, :-1])  # (Nw, Nt)
	weights = phases * dts[None, :]  # (Nw, Nt)
	F_trans = jax.tree_util.Partial(jnp.einsum, "wt,tde->wde", weights)

	def if_trans(U_t, ori_H):
		if_H = U_t.T.conj() @ ori_H @ U_t
		if_H_t = if_H - jnp.trace(P @ if_H) / Tr_P * I
		return if_H_t

	if_trans_vc = jax.vmap(if_trans, (0, None), 0)
	if_trans_vv = jax.vmap(if_trans, (0, 0), 0)

	def func_robust_infide_from_fw(if_H_ws, S_H_ws):
		F_H_ws = jnp.einsum("wde,wde->w", P @ if_H_ws, jnp.conj(if_H_ws)) / Tr_P
		# robust_infide = jnp.real((F_H_ws * S_H_ws)[:-1] @ dws / jnp.pi)
		return jnp.real(jsci_trapezoid(F_H_ws * S_H_ws, ws)) / jnp.pi

	def func_robust_infide_Vn(Vn, S_Vn_ws):
		if_Vn_ts = if_trans_vc(U_ts, Vn)  # (Nt, dim, dim)
		if_Vn_ws = F_trans(if_Vn_ts)
		return func_robust_infide_from_fw(if_Vn_ws, S_Vn_ws)

	func_robust_infide_Vns = jax.vmap(func_robust_infide_Vn, (0, 0), 0)

	robust_infides = jnp.zeros(N_Vns + 1)
	robust_infides = robust_infides.at[:N_Vns].set(func_robust_infide_Vns(Vns, S_Vns_ws))

	if_Hc_ts = if_trans_vv(U_ts, Hc_ts)  # (Nt, dim, dim)
	if_Hc_ws = F_trans(if_Hc_ts)
	robust_infide_Hc = func_robust_infide_from_fw(if_Hc_ws, S_Hc_ws)
	robust_infides = robust_infides.at[N_Vns].set(robust_infide_Hc)
	return UT, robust_infides

# def calcu_dv_loss(
# 		ucs_ts,
# 		dts
# ):
# 	return (
# 			jnp.sum((jnp.diff(ucs_ts, axis=1) / dts[None, :-1]) ** 2)
# 			+ jnp.sum((ucs_ts[:, 0] / dts[0]) ** 2)
# 			+ jnp.sum((ucs_ts[:, -1] / dts[-1]) ** 2)
# 	)
